package cn.redragon.soa.common.util.jwt;

import static org.springframework.util.DigestUtils.appendMd5DigestAsHex;
import static org.springframework.util.DigestUtils.md5DigestAsHex;

import io.jsonwebtoken.Claims;
import io.jsonwebtoken.ExpiredJwtException;
import io.jsonwebtoken.JwtException;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.MalformedJwtException;
import java.security.KeyFactory;
import java.security.PublicKey;
import java.security.spec.X509EncodedKeySpec;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import lombok.Getter;
import lombok.SneakyThrows;
import org.apache.commons.codec.binary.Base64;

public class RsaJwtTokenVerifier {

    @Getter
    private final PublicKey publicKey;

    @SneakyThrows
    public RsaJwtTokenVerifier(final String publicKeyString) {
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        byte[] encodedPb = Base64.decodeBase64(publicKeyString);
        X509EncodedKeySpec keySpecPb = new X509EncodedKeySpec(encodedPb);

        this.publicKey = keyFactory.generatePublic(keySpecPb);
    }

    /**
     * Parse body of Jwt and get the subject
     *
     * @throws MalformedJwtException if the token string is not a valid JWS
     * @throws ExpiredJwtException   if the token's expiration time before the current time
     */

    /**
     * Parse body of Jwt and get the roles if exists
     *
     * @return Set of roles, or emptyList if 'rol' not exists in token
     * @throws MalformedJwtException if the token string is not a valid JWS
     * @throws ExpiredJwtException   if the token's expiration time before the current time
     */
    public Collection<String> getUserRole(String token) {
        Claims body = getTokenBody(token);
        return new ArrayList<>();
    }


    /**
     * Validate whether the token is expired
     *
     * @throws JwtException if the token string is not a valid JWS
     */
    public boolean isTokenExpired(String token) {
        try {
            return getTokenBody(token).getExpiration().before(new Date());
        } catch (JwtException | IllegalArgumentException e) {
            return true;
        }
    }

    /**
     * Validate whether the token is expired
     *
     * @throws MalformedJwtException if the token string is not a valid JWS
     * @throws ExpiredJwtException   if the token's expiration time before the current time
     */
    public Date getExpiration(String token) {
        return getTokenBody(token).getExpiration();
    }

    /**
     * Get specified claim from jwt body
     *
     * @throws MalformedJwtException if the token string is not a valid JWS
     * @throws ExpiredJwtException   if the token's expiration time before the current time
     */
    public <T> T getClaim(String token, String claim, Class<T> clazz) {
        Claims body = getTokenBody(token);

        return body.get(claim, clazz);
    }

    /**
     * Parse the token with secret key, return the body of Jwt
     *
     * @throws MalformedJwtException if the token string is not a valid JWS
     * @throws ExpiredJwtException   if the token's expiration time before the current time
     */
    public Claims getTokenBody(String token) {
        return Jwts.parserBuilder()
            .setSigningKey(publicKey)
            .build()
            .parseClaimsJws(token)
            .getBody();
    }

    /**
     * In case public key is changed, we need to know
     *
     * @param password password
     * @return hashed value
     */
    public String appendPublicKeyHash(String password) {
        return appendMd5DigestAsHex(this.publicKey.getEncoded(), new StringBuilder(md5DigestAsHex(password.getBytes()))).toString();
    }
}
